import os
import json
import pandas as pd
from langchain_openai import ChatOpenAI
from dotenv import load_dotenv
from sklearn.metrics import classification_report
from tenacity import retry, stop_after_attempt, wait_fixed
load_dotenv()
# 步骤1：通过os.getenv读取class_file文件，创建 id2name 字典 与 name2id字典
id2name = {i: line.strip() for i, line in enumerate(open(os.getenv("class_file"), encoding="utf-8"))}
name2id = {v: k for k, v in id2name.items()}


# 步骤2：deepseek-chat 初始化语言模型 ，返回值约束model_kwargs={"response_format": {"type": "json_object"}}
#ChatOpenAI 是 langchain_openai 模块中的一个类，用于与支持模型交互（如 DeepSeek 的 deepseek-chat 模型）。
llm = ChatOpenAI(
    base_url=os.getenv("base_url"),
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    model="deepseek-chat",
    model_kwargs={"response_format": {"type": "json_object"}}
)

# 步骤 3：定义带重试机制的 LLM 调用函数
#"""调用 LLM，若失败则重试最多 3 次，每次间隔 2 秒"""
#@retry(stop=stop_after_attempt(3), wait=wait_fixed(2))
@retry(stop=stop_after_attempt(3), wait=wait_fixed(2))
def invoke_llm(prompt):
    """调用 LLM，若失败则重试最多 3 次，每次间隔 2 秒"""
    return llm.invoke(prompt)


# 步骤 4：定义读取数据文件read_data
def read_data(file_path):
    """读取 tab 分隔的新闻标题和标签数据"""
    df = pd.read_csv(file_path, sep='\t', header=None, names=['title', 'label'], encoding='utf-8')
    return df['title'].tolist(), df['label'].tolist()


# 步骤 5：分类新闻标题
def classify_news(title: str) -> dict:
    """将新闻标题分类到指定类别，返回分类结果"""
    # 定义提示词，包含类别说明和示例
    prompt = [
        {
            "role": "system",
            "content": """
你是一名新闻分类审核员，任务是将新闻标题分类到以下类别之一：
finance, realty, stocks, education, science, society, politics, sports, game, entertainment。
请根据标题内容和以下关键词与示例，匹配最相关的类别。如果标题涉及教育机构但核心是社会贡献，优先归为 society。
返回 JSON 格式：{"category": "类别", "reason": "分类原因"}

类别关键词与示例：
- finance: 银行、信用卡、贷款、利率 (例: "各银行信用卡挂失费迥异")
- realty: 房产、地价、楼盘 (例: "东5环海棠公社230平准现房")
- stocks: 股市、股指、期货 (例: "金证顾问：过山车行情意味着什么")
- education: 学校、考试、招生 (例: "中华女子学院仅1专业招男生")
- science: 技术、网站、宇航 (例: "“手机钱包”亮相科博会")
- society: 社会事件、犯罪、公益 (例: "82岁老太为学生做饭扫地44年")
- politics: 政策、国际关系 (例: "查韦斯称愿为俄罗斯提供空军基地")
- sports: 比赛、运动员、奥运 (例: "卡佩罗：德国脚生猛的原因")
- game: 电子游戏、网游、电竞 (例: "《赤壁OL》攻城战硝烟又起")
- entertainment: 明星、影视、综艺 (例: "冯德伦徐若瑄隔空传情")
"""
        },
        {
            "role": "user",
            "content": f"新闻标题：'{title}'，请分类并说明原因。"
        }
    ]

    #  调用 大模型获取分类结果 invoke_llm(prompt)
    response = invoke_llm(prompt)
    # 获取 response.content大模型返回回答内容，将json字符串转为字典
    result = json.loads(response.content)
    # 返回分类结果，包含类别（category）和原因（reason）
    return {
        "category": result.get("category", "society"),  # 默认 society
        "reason": result.get("reason", "未明确分类，归为社会类别")
    }


# 步骤 6：大模型预测逻辑调用、模型验证及评估
if __name__ == '__main__':
    #定义要预测的文件路径
    file_path = "../01-data/dev2.txt"
    #  read_data读取数据
    titles, true_labels = read_data(file_path)

    # 定义预测结果pred_labels以及最终结果保存results
    pred_labels = []
    results = []

    # 对每个标题进行分类
    for title in titles:
        print(f"处理标题: {title}")
        # classify_news调用大模型进行预测，得到result
        result = classify_news(title)
        # 获取result中的预测类别
        category = result["category"]
        # 通过name2id将预测类别转换为数字标签
        pred_label = name2id.get(category, 5)  # 默认 society 为 5
        #  保存预测类别pred_labels、以及结果results（包含title、category、reason）
        pred_labels.append(pred_label)
        results.append({
            "title": title,
            "category": category,
            "reason": result["reason"]
        })

    # 遍历results打印分类结果
    print("\n分类结果：")
    for result in results:
        print(result)

    # 评估性能
    print("\n评估指标：")
    report = classification_report(
        true_labels,
        pred_labels,
    )
    print(report)